#include "blockwiseIntegrated.hpp"
#include "bi_calcBlocksize.hpp"
#include "solverBase.hpp"
#include "sweSolver.hpp"
#include "meshTypeFunctions.hpp"
#include "structMesh_testFunctions.hpp"
#include "bc_base_all.hpp"
#include "bc_swe_all.hpp"
#include <catch2/catch_test_macros.hpp>
#include <catch2/catch_approx.hpp>

using Catch::Approx;
using namespace hms;

TEST_CASE("Block size calculation"){
	using namespace hms::swe::block;
	#define HMS_REQ_NN(MESH_SIZE, APPROX_SIZE, RES) \
	REQUIRE( calculateSize_nextNeighbour(MESH_SIZE, APPROX_SIZE) == RES )

	HMS_REQ_NN(10, 64, 10);
	HMS_REQ_NN(63, 64, 63);
	HMS_REQ_NN(65, 64, 32);
	HMS_REQ_NN(513, 64, 57);
	HMS_REQ_NN(514, 64, 57);
	HMS_REQ_NN(5000, 64, 63);

	#define HMS_REQ_TH(MESH_SIZE, APPROX_SIZE, RES) \
	REQUIRE( calculateSize_threshold(MESH_SIZE, APPROX_SIZE, 0.875) == RES )

	HMS_REQ_TH(10, 64, 64);
	HMS_REQ_TH(63, 64, 63);
	HMS_REQ_TH(65, 64, 56);
	HMS_REQ_TH(513, 64, 57);
	HMS_REQ_TH(514, 64, 58);
	HMS_REQ_TH(5000, 64, 61);

	#undef HMS_REQ_NN
	#undef HMS_REQ_TH
}

TEST_CASE("Blockwise index helper"){
	auto mesh { uniMeshAs<UniMesh>(12,13,1,1,false) };
	Index
		blocksizeX {3},
		blocksizeY {4};
	
	swe::IndexHelper ih { swe::indexHelper(blocksizeX, blocksizeY, mesh) };

	const auto& m  { ih.mesh };
	REQUIRE(m.nx == 12);
	REQUIRE(m.ny == 13);
	REQUIRE(m.ng ==  1);
	REQUIRE(m.nxg == 14);
	REQUIRE(m.nyg == 15);
	REQUIRE(m.nOuterLayers == 2);
	REQUIRE(m.nxInner == 8);
	REQUIRE(m.nyInner == 9);
	// REQUIRE(m.begOuter == 15);
	// REQUIRE(m.begInner == 45);

	const auto& ib { ih.blocks };
	REQUIRE(ib.sizeX == blocksizeX);
	REQUIRE(ib.sizeY == blocksizeY);
	REQUIRE(ib.nFullX == 2);
	REQUIRE(ib.nFullY == 2);
	REQUIRE(ib.remSizeX  == 2);
	REQUIRE(ib.remSizeY  == 1);

	// REQUIRE(ib.remRight  == 51);
	// REQUIRE(ib.remTop    == 157);
	// REQUIRE(ib.remCorner == 163);
	
	// const auto& ob { ih.outer };
	// REQUIRE(ob.bottom ==  17);
	// REQUIRE(ob.top    == 171);
	// REQUIRE(ob.left   ==  43);
	// REQUIRE(ob.right  ==  53);
	// REQUIRE(ob.bottomRem ==  23);
	// REQUIRE(ob.topRem    == 177);
	// REQUIRE(ob.leftRem   == 155);
	// REQUIRE(ob.rightRem  == 165);

	// const auto& c  { ih.corners };
	// REQUIRE(c.bottomLeft  ==  15);
	// REQUIRE(c.bottomRight ==  25);
	// REQUIRE(c.topLeft     == 169);
	// REQUIRE(c.topRight    == 179);

	const auto& beg { ih.beg };
	REQUIRE(beg.mesh   ==  1);
	REQUIRE(beg.blocks ==  3);
	REQUIRE(beg.remX   ==  9);
	REQUIRE(beg.remY   == 11);
	REQUIRE(beg.outerX == 11);
	REQUIRE(beg.outerY == 12);
}

/* only BlockLocation::inner is currently being tested here! */
TEST_CASE( "Block creation and corresponding variables" ){
	using MeshType = UniMesh;
	using namespace hms::swe;
	using B = BlockLocation;

	Index
		nx {17},
		ny {12},
		ng {1},
		blocksizeX {4},
		blocksizeY {3};
	
	auto mesh { uniMeshAs<MeshType>(nx, ny, 1, 1) };
	auto ih { indexHelper(blocksizeX, blocksizeY, mesh) };

	Block<B::inner> block { ih, 7, 3, 4, 3 };
	REQUIRE( (block.firstCell() == Array2i{7,3}).all() );
	REQUIRE( (block.size() == Array2i{4,3}).all() );
	REQUIRE( (block.sizeRead() == Array2i{8,7}).all() );
	REQUIRE( block.nRead() == 56 );
	REQUIRE( block.nWrite() == 12 );
	REQUIRE( block.xBeg() == 7 );
	REQUIRE( block.yBeg() == 3 );
	REQUIRE( block.nx() == 4 );
	REQUIRE( block.ny() == 3 );
	REQUIRE( block.nxRead() == 8 );
	REQUIRE( block.nyRead() == 7 );

	Array3Xs mem ( 3, block.nRead() );
	auto mem_map { block.map(mem) };
	REQUIRE( mem_map.rows() == 3 );
	REQUIRE( mem_map.cols() == block.nRead() );

	ArrayXXs arr ( 2, nFieldCols(mesh) );
	arr.reshaped().setLinSpaced( 0, arr.rows()*nFieldCols(mesh)-1 );

	auto rb { block.template readBlock<2>(arr) };
	REQUIRE( rb.rows() == 8*arr.rows() );
	REQUIRE( rb.cols() == 7 );
	REQUIRE( ( rb.col(0) == ArrayNs<16>::LinSpaced(48, 63) ).all() );

	auto wb { block.template writeBlock<2>(arr) };
	REQUIRE( wb.rows() == 4*arr.rows() );
	REQUIRE( wb.cols() == 3 );
	REQUIRE( ( wb.col(0) == ArrayNs<8>::LinSpaced(128, 135) ).all() );

	mem_map.reshaped().setLinSpaced( 0, mem_map.size()-1 );
	auto wb_mem { block.writeBlock_mem(mem_map) };
	REQUIRE( wb_mem.rows() == 4*mem.rows() );
	REQUIRE( wb_mem.cols() == wb.cols() );
	REQUIRE( ( wb_mem.col(0) == ArrayNs<12>::LinSpaced(54, 65) ).all() );

	/* mesh rows are represented by columns in eigen */
	auto rowMap { block.template rowMap<2>(wb, 1) };
	REQUIRE( rowMap.size() == wb.rows() );
	REQUIRE( ( rowMap.reshaped() == wb.col(1) ).all() );

	auto rowMapRead { block.template rowMapRead<2>(rb, 1) };
	REQUIRE( rowMapRead.size() == rb.rows() );
	REQUIRE( ( rowMapRead.reshaped() == rb.col(1) ).all() );

	REQUIRE( block.hBeg() == 40 );
	REQUIRE( block.vBeg() == 221 + 42 );
	REQUIRE( (block.lrBeg_block<EdgeRange::horizontal>() == Array2i{10,18}).all() );
	REQUIRE( (block.lrBeg_block<EdgeRange::vertical  >() == Array2i{17,18}).all() );
	REQUIRE( (block.lrBeg_mesh <EdgeRange::horizontal>() == Array2i{45,64}).all() );
	REQUIRE( (block.lrBeg_mesh <EdgeRange::vertical  >() == Array2i{63,64}).all() );
}

TEST_CASE( "SolverThread functions" ){
	// using MeshType = UniMesh;
	using MeshType = RectMesh;
	using namespace hms::swe;
	using B = BlockLocation;

	Index
		nx {17},
		ny {12},
		ng {1},
		blocksizeX {4},
		blocksizeY {3};
	scalar
		dx {0.2},
		dy {0.1},
		maxCourant {0.5},
		d {2},
		vx {3},
		vy {4},
		dt { 0.1 };
	Array2s v {vx, vy};
	
	auto mesh { uniMeshAs<MeshType>(nx, ny, dx, dy, ng) };
	auto ih { indexHelper(blocksizeX, blocksizeY, mesh) };
	Time time {10};
	time.step( dt );

	Block<B::inner> block { ih, 7, 3, 4, 3 };

	SWESolver solver {mesh, time};
	solver.settings().maxCourant = maxCourant;

	Index nCols { nFieldCols(mesh) };
	ArrayXXs
		dNew { 1, nCols },
		vNew { 2, nCols };
	dNew = d;
	vNew.colwise() = v;

	SolverThread thread { solver, dNew, vNew, ih };
	auto dBlock { block.template writeBlock<1>(dNew) };
	auto vBlock { block.template writeBlock<2>(vNew) };

	Array3XsMap q { block.map( thread.qBuf   () ) };
	Array3XsMap f { block.map( thread.fluxBuf() ) };

	SECTION("Time step"){
		scalar ds { 4 * dx * dy / ( 2*(dx+dy) ) };
		scalar vmax { Vector2s{vx, vy}.norm() + std::sqrt(phys::g*d) };
		REQUIRE(
			thread.timestep(block, dBlock, vBlock) ==
			Approx( maxCourant * ds / vmax )
		);
	}
	SECTION("Mass sources"){
		scalar rain { 0.1 };

		solver.useRainfall(true);
		solver.r() = SourceField<Uniform, MeshType>{ mesh, Array1s{rain} };
		q = 0;
		thread.addMassSources(block, q);
		auto massStorage { block.writeBlock_mem(q)
			.reshaped( 3, block.nWrite() )
			.template topRows<1>()
		};
		REQUIRE( massStorage.isApproxToConstant( dt * rain) );
	}
	SECTION("Friction"){
		Array3s flux { d, 2, 0 };
		f.colwise() = flux;

		scalar
			vNorm { v.matrix().norm() },
			manning { 0.025 },
			cf { phys::g * std::pow(manning, 2) / std::pow(2., 1./3) },
			fr { dt * cf * vNorm },
			modCf { cf * dt / d };

		auto applyFriction = [&]( scalar fluxComponent, scalar vComponent )
			-> scalar
		{
			return
				( fluxComponent - fr * vComponent ) /
				( 1 + modCf * (vNorm + std::pow(vComponent, 2) / vNorm) );
		};

		Array3s modFlux {
			flux[0],
			applyFriction( flux[1], vx ),
			applyFriction( flux[2], vy )
		};

		thread.computeFriction(block, f, dBlock, vBlock);
		REQUIRE(
			block.writeBlock_mem(f).reshaped( 3, block.nWrite() )
			.isApprox(
			modFlux.replicate( 1, block.nWrite() )
			)
		);
	}
}

// TEST_CASE( "Blockwise output test", "[pureOutput]" ){
// 	// using MeshType = Mesh;
// 	using MeshType = UniMesh;
// 	// using MeshType = RectMesh;
// 	// using MeshType = StructMesh;

// 	Index
// 		nx {17},
// 		ny {12},
// 		ng {1},
// 		blocksizeX {4},
// 		blocksizeY {3};

// 	auto mesh { uniMeshAs<MeshType>(nx, ny, 1, 1, true) };

// 	Time time {10};
// 	time.writeAtTimeIntervals(false);
// 	time.writeStepInterval(1);

// 	SWESolver<MeshType> solver {mesh, time};
	
// 	solver.v().values().colwise() = Array2s{1, 0};
// 	solver.d().values() = 1;
// 	solver.z().values() = 0;

// 	for ( Field<MeshType>& field : solver.fields() ){
// 		for ( const BoundaryPatch& patch : mesh.boundaryPatches() ){
// 			field.ghostValues(patch) = 99;
// 			field.setBoundaryCondition( StaticValue{field, patch, false );
// 		}
// 	}

// 	solver.useRainfall(false);
// 	solver.initialise();

// 	omp_set_num_threads(1);

	// swe::integratedTimeLoop(solver, time, blocksizeX, blocksizeY,
	// 	[](){}, [&](){
	// 		std::cout
	// 			<< "Time: " << time.current()
	// 			<< ", step: " << time.step()
	// 			<< ", nSteps: " << time.nSteps()
	// 			<< "\nd:\n";
	// 		// printAsStructMesh(solver.d().values(), mesh);
	// 	},
	// 	false,
	// 	false,
	// 	3
	// );
// }